import psycopg2
from psycopg2.extras import RealDictCursor
from typing import Optional, List, Dict, Any
import logging
from config import config

class WebStore:
    """Web数据存储类，用于连接PostgreSQL数据库并进行数据操作"""
    
    def __init__(self):
        self.connection = None
        self.logger = logging.getLogger(__name__)
        
    def connect(self) -> bool:
        """连接PostgreSQL数据库"""
        try:
            # 构建连接字符串
            connection_params = {
                'host': config.POSTGRES_HOST,
                'port': int(config.POSTGRES_PORT),
                'database': config.POSTGRES_DB,
                'user': config.POSTGRES_USER,
                'password': config.POSTGRES_PASSWORD
            }
            
            self.connection = psycopg2.connect(**connection_params)
            self.connection.autocommit = False
            self.logger.info("成功连接到PostgreSQL数据库")
            return True
            
        except Exception as e:
            self.logger.error(f"连接数据库失败: {e}")
            return False
    
    def disconnect(self):
        """断开数据库连接"""
        if self.connection:
            self.connection.close()
            self.connection = None
            self.logger.info("已断开数据库连接")
    
    def __enter__(self):
        """上下文管理器入口"""
        if not self.connect():
            raise Exception("无法连接到数据库")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器退出"""
        self.disconnect()
    
    def insert_company_announcement(self, stock_code: str, company_name: str, 
                                  announcement_title: str, announcement_date: str,
                                  exchange: Optional[str] = None) -> Optional[int]:
        """
        插入公司公告主信息（如果不存在的话）
        
        Args:
            stock_code: 股票代码
            company_name: 公司名称
            announcement_title: 公告标题
            announcement_date: 公告日期
            exchange: 交易所标识（可选）
            
        Returns:
            插入或已存在记录的ID，失败返回None
        """
        if not self.connection:
            self.logger.error("数据库未连接")
            return None
            
        try:
            cursor = self.connection.cursor()
            
            # 先检查记录是否已存在（考虑 exchange 字段）
            if exchange is None:
                check_sql = """
                SELECT id FROM company_announcements 
                WHERE stock_code = %s AND announcement_title = %s AND announcement_date = %s
                  AND exchange IS NULL;
                """
                cursor.execute(check_sql, (stock_code, announcement_title, announcement_date))
            else:
                check_sql = """
                SELECT id FROM company_announcements 
                WHERE stock_code = %s AND announcement_title = %s AND announcement_date = %s
                  AND exchange = %s;
                """
                cursor.execute(check_sql, (stock_code, announcement_title, announcement_date, exchange))
            existing_record = cursor.fetchone()
            
            if existing_record:
                announcement_id = existing_record[0]
                self.logger.info(f"公告已存在，ID: {announcement_id}")
                cursor.close()
                return announcement_id
            
            # 记录不存在，插入新记录（包含 exchange 字段）
            insert_sql = """
            INSERT INTO company_announcements (stock_code, company_name, announcement_title, announcement_date, exchange)
            VALUES (%s, %s, %s, %s, %s)
            RETURNING id;
            """
            
            cursor.execute(insert_sql, (stock_code, company_name, announcement_title, announcement_date, exchange))
            result = cursor.fetchone()
            
            self.connection.commit()
            cursor.close()
            
            announcement_id = result[0] if result else None
            self.logger.info(f"成功插入新公告数据，ID: {announcement_id}")
            return announcement_id
            
        except Exception as e:
            self.logger.error(f"插入公告数据失败: {e}")
            if self.connection:
                self.connection.rollback()
            return None
    
    def insert_announcement_link(self, announcement_id: int, href: str, link_text: str,
                               is_pdf: bool = True, original_href: str = None, 
                               link_order: int = 0) -> bool:
        """
        插入公告链接信息（如果不存在的话）
        
        Args:
            announcement_id: 公告ID
            href: 完整链接地址
            link_text: 链接文本
            is_pdf: 是否为PDF链接
            original_href: 原始链接地址
            link_order: 链接顺序
            
        Returns:
            插入是否成功（包括已存在的情况）
        """
        if not self.connection:
            self.logger.error("数据库未连接")
            return False
            
        try:
            cursor = self.connection.cursor()
            
            # 先检查链接是否已存在
            check_sql = """
            SELECT id FROM announcement_links 
            WHERE announcement_id = %s AND href = %s;
            """
            
            cursor.execute(check_sql, (announcement_id, href))
            existing_link = cursor.fetchone()
            
            if existing_link:
                self.logger.info(f"链接已存在，跳过插入: {href}")
                cursor.close()
                return True
            
            # 链接不存在，插入新链接
            insert_sql = """
            INSERT INTO announcement_links (announcement_id, href, link_text, is_pdf, original_href, link_order)
            VALUES (%s, %s, %s, %s, %s, %s);
            """
            
            cursor.execute(insert_sql, (announcement_id, href, link_text, is_pdf, original_href, link_order))
            
            self.connection.commit()
            cursor.close()
            
            self.logger.info(f"成功插入新链接数据，公告ID: {announcement_id}")
            return True
            
        except Exception as e:
            self.logger.error(f"插入链接数据失败: {e}")
            if self.connection:
                self.connection.rollback()
            return False
    
    def insert_announcement_with_links(self, stock_code: str, company_name: str,
                                     announcement_title: str, announcement_date: str,
                                     links: List[Dict[str, Any]],
                                     exchange: Optional[str] = None) -> Optional[int]:
        """
        插入公告及其相关链接（事务操作）
        
        Args:
            stock_code: 股票代码
            company_name: 公司名称  
            announcement_title: 公告标题
            announcement_date: 公告日期
            links: 链接列表，每个元素包含href, link_text, is_pdf, original_href, link_order
            exchange: 交易所标识（可选）
            
        Returns:
            插入的公告ID，失败返回None
        """
        if not self.connection:
            self.logger.error("数据库未连接")
            return None
            
        try:
            # 插入主公告信息
            announcement_id = self.insert_company_announcement(
                stock_code, company_name, announcement_title, announcement_date, exchange
            )
            
            if not announcement_id:
                return None
            
            # 插入链接信息
            for i, link in enumerate(links):
                link_text = link.get('text', '').strip()
                if not link_text:  # 跳过空链接文本
                    continue
                
                success = self.insert_announcement_link(
                    announcement_id=announcement_id,
                    href=link.get('href', ''),
                    link_text=company_name + ":" + link_text,
                    is_pdf=link.get('is_pdf', True),
                    original_href=link.get('original_href'),
                    link_order=link.get('link_order', i)
                )
                
                if not success:
                    self.logger.error(f"插入链接失败，回滚事务")
                    self.connection.rollback()
                    return None
            
            self.logger.info(f"成功插入公告及链接数据，公告ID: {announcement_id}")
            return announcement_id
            
        except Exception as e:
            self.logger.error(f"插入公告及链接数据失败: {e}")
            if self.connection:
                self.connection.rollback()
            return None
    
    def get_announcement_by_id(self, announcement_id: int) -> Optional[Dict[str, Any]]:
        """根据ID查询公告信息"""
        if not self.connection:
            return None
            
        try:
            cursor = self.connection.cursor(cursor_factory=RealDictCursor)
            
            sql = """
            SELECT * FROM company_announcements WHERE id = %s;
            """
            
            cursor.execute(sql, (announcement_id,))
            result = cursor.fetchone()
            cursor.close()
            
            return dict(result) if result else None
            
        except Exception as e:
            self.logger.error(f"查询公告数据失败: {e}")
            return None